package loadbalance;

import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

import static java.util.Objects.hash;

/**
 * 基于 hash 一致性的负载均衡
 *
 * @author Daniel
 */
public class ConsistentHashLoadBalance {

    private ConsistentHashSelector<?> selector;

    /**
     * 虚拟节点数
     */
    private final int replicaNumber;

    public ConsistentHashLoadBalance() {
        this.replicaNumber = 160;
    }

    public ConsistentHashLoadBalance(int replicaNumber) {
        this.replicaNumber = replicaNumber;
    }

    public Invoker doSelect(List<Invoker> invokers, String key) {

        if (selector == null) {
            // 创建新的 ConsistentHashSelector
            selector = new ConsistentHashSelector(invokers, replicaNumber);
        }

        // 调用 ConsistentHashSelector 的 select 方法选择 Invoker
        return selector.select(key);
    }

    private static final class ConsistentHashSelector<T> {

        /**
         * 使用 TreeMap 存储 Invoker 虚拟节点
         */
        private final TreeMap<Long, Invoker> virtualInvokers;

        /**
         * 虚拟节点数
         */
        private final int replicaNumber;

        ConsistentHashSelector(List<Invoker> invokers, int replicaNumber) {
            this.virtualInvokers = new TreeMap<>();
            this.replicaNumber = replicaNumber;
            for (Invoker invoker : invokers) {
                int id = invoker.getId();
                for (int i = 0; i < replicaNumber / 4; i++) {
                    // 对 address + i 进行 md5 运算，得到一个长度为16的字节数组
                    byte[] digest = md5("" + id + i);
                    // 对 digest 部分字节进行4次 hash 运算，得到四个不同的 long 型正整数
                    for (int h = 0; h < 4; h++) {
                        // h = 0 时，取 digest 中下标为 0 ~ 3 的4个字节进行位运算
                        // h = 1 时，取 digest 中下标为 4 ~ 7 的4个字节进行位运算
                        // h = 2, h = 3 时过程同上
                        long m = hash(digest, h);
                        // 将 hash 到 invoker 的映射关系存储到 virtualInvokers 中，
                        // virtualInvokers 需要提供高效的查询操作，因此选用 TreeMap 作为存储结构
                        virtualInvokers.put(m, invoker);
                    }
                }
            }
        }

        private byte[] md5(String s) {
            MessageDigest md5 = null;
            try {
                md5 = MessageDigest.getInstance("MD5");
            } catch (NoSuchAlgorithmException e) {
                // nothing
            }
            assert md5 != null;
            md5.update(s.getBytes());
            return md5.digest();
        }

        public Invoker select(String key) {
            // 对参数 key 进行 md5 运算
            byte[] digest = md5(key);
            // 取 digest 数组的前四个字节进行 hash 运算，再将 hash 值传给 selectForKey 方法，
            // 寻找合适的 Invoker
            return selectForKey(hash(digest, 0));
        }

        private Invoker selectForKey(long hash) {
            // 到 TreeMap 中查找第一个节点值大于或等于当前 hash 的 Invoker
            Map.Entry<Long, Invoker> entry = virtualInvokers.tailMap(hash, true).firstEntry();
            // 如果 hash 大于 Invoker 在圆环上最大的位置，此时 entry = null，
            // 需要将 TreeMap 的头节点赋值给 entry
            if (entry == null) {
                entry = virtualInvokers.firstEntry();
            }

            // 返回 Invoker
            return entry.getValue();
        }
    }

}
